﻿/**
 * Created by zywy on 2023-06-06
 * All rights reserved.
 */

#pragma once
#ifndef _MIDDLEWARE_NUMAMEM_
#define _MIDDLEWARE_NUMAMEM_

#if !defined(__CYGWIN__) && ENABLE_NUMA

#include <numa.h>
#include <pthread.h>
#include <memory>
#include <map>
#include <functional>
#include <dlfcn.h>
#include "logger.h"

#define NumaMemI NumaMem::shareInstance()

// bind_cpu()的封装，第2个参数传入用户代码块
#define numamem_bind_cpu(cpu, ...)                      \
    NumaMemI->unbind(NumaMemI->bind_cpu(cpu), [&]() {   \
        __VA_ARGS__                                     \
    });

// bind_node()的封装，第2个参数传入用户代码块
#define numamem_bind_node(node, ...)        \
    NumaMemI->unbind(NumaMemI->bind_node(node), [&]() {   \
        __VA_ARGS__                                     \
    });

class NumaMem {
public:
    static std::shared_ptr<NumaMem> shareInstance() {
        static std::shared_ptr<NumaMem> NumaMemInstance = std::make_shared<NumaMem>();
        return NumaMemInstance;
    };

    NumaMem();
    ~NumaMem();

    /**
     * 只在指定的node上进行内存分配,
     * 直到下次调用 bind / unbind 前有效
     * 
     * @param node 
     * @return 
     *     id on 成功，返回一个id，用于unbind
     *     -1 on 无效的node
     *     -2 on NUMA功能不可用
     */
    int bind_node(int node);

    /**
     * 只在指定的cpu所在的node上进行内存分配, 
     * 直到下次调用 bind / unbind 前有效
     * 
     * @param cpu 
     * @return 
     *     id on 成功，返回一个id，用于unbind
     *     -1 on 无效的cpu
     *     -2 on NUMA功能不可用
     */
    int bind_cpu(int cpu);

    /**
     * 还原内存分配策略
     * 
     * @param id 由 bind_cpu / bind_node 返回
     * @return 
     *      0 on 成功
     *     -1 on 无效的id
     *     -2 on NUMA功能不可用
     */
    int unbind(int id);

    int unbind(int id, std::function<void(void)> f) {
        f();
        return unbind(id);
    }

    /**
     * 在指定的cpu所在的node上进行new操作
     * 
     * @return
     *      new创建的对象的指针，NUMA功能不可用或者cpu无效时使用普通的new
     */
    template <typename T, int cpu, class ... Args>
    T *new_cpu(Args&& ... args) {
        int id;

        if (lib_handle == nullptr || dl_numa_available() < 0) {
            return new T(args...);
        }

        if (cpu < 0 || cpu >= dl_numa_num_configured_cpus()) {
            LOG_WARN("NumaMem::new_cpu: Invalid cpu %d", cpu);
            return new T(args...);
        }

        id = bind_cpu(cpu);
        T *obj = new T(args...);
        unbind(id);
        return obj;
    }

    /**
     * 在指定的node上进行new操作
     * 
     * @return
     *      new创建的对象的指针，NUMA功能不可用或者cpu无效时使用普通的new
     */
    template <typename T, int node, class ... Args>
    T *new_node(Args&& ... args) {
        int id;

        if (lib_handle == nullptr || dl_numa_available() < 0) {
            return new T(args...);
        }

        if (node < 0 || node >= dl_numa_num_configured_nodes()) {
            LOG_WARN("NumaMem::new_node: Invalid node %d", node);
            return new T(args...);
        }

        id = bind_node(node);
        T *obj = new T(args...);
        unbind(id);
        return obj;
    }

    std::function<void(struct bitmask *&)> bitmask_del = [&](struct bitmask *&mask) {
        dl_numa_bitmask_free(mask);
    };

    // dynamic-load symbols
    int(*dl_numa_available)();
    struct bitmask *(*dl_numa_bitmask_alloc)(unsigned int);
    void (*dl_numa_bitmask_free)(struct bitmask *);
    struct bitmask *(*dl_numa_bitmask_setbit)(struct bitmask *, unsigned int);
    int (*dl_numa_num_configured_nodes)();
    int (*dl_numa_num_configured_cpus)();
    struct bitmask *(*dl_numa_get_membind)(void);
    void (*dl_numa_set_membind)(struct bitmask *nodemask);
    int (*dl_numa_node_of_cpu)(int cpu);

private:
    void *lib_handle = nullptr;
    pthread_spinlock_t lock;
    struct bitmask *mask;
    int old_membind_policy_id = 0;
    std::map<int, std::shared_ptr<struct bitmask>> old_membind_policy;
    std::map<std::string, void *> dlfuncs;
};

#endif
#endif
